#
# Copyright (c) 2011, Michael Smith
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#     o  Redistributions of source code must retain the above copyright
#        notice, this list of conditions and the following disclaimer.
#     o  Redistributions in binary form must reproduce the above copyright
#        notice, this list of conditions and the following disclaimer in the
#        documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#

from struct import pack, unpack
import hmac
import md4
import time
import random
import math
try:
	from hashlib import md5
except ImportError:
	from md5 import md5
	
"""
	The NT LAN Manager (NTLM) Authentication Protocol is used in Microsoft Windows for authentication
	between clients and servers. NTLM is a challenge-response style authentication protocol.
	
	CLIENT                                     SERVER
	  [<---------Application Message(s)--------->]
	  [--------------NTLM.Negotiate------------->]
	  [<-------------NTLM.Challenge--------------]
	  [-------------NTLM.Authenticate----------->]
	  [<---------Application Messages(s)-------->]
"""

# NTLMSSP Signature string
NTLMSSP_SIGNATURE = 'NTLMSSP\x00'
VERSION = 0x0f0000001db00106 # Windows 7, build 7600, NTLM revision 15

# Security Protocol OID strings
NTLMSSP_OID = '\x2b\x06\x01\x04\x01\x82\x37\x02\x02\x0a'
SPNEGO_OID = '\x2b\x06\x01\x05\x05\x02'

#NTLMSSP Flags
NTLMSSP_NEGOTIATE_UNICODE					= 0x00000001  # Unicode strings are supported for use in security buffer data. 
NTLM_NEGOTIATE_OEM							= 0x00000002  # OEM strings are supported for use in security buffer data. 
NTLMSSP_REQUEST_TARGET						= 0x00000004  # Requests server include authentication realm in the challenge message. 
NTLMSSP_NEGOTIATE_SIGN						= 0x00000010  # Messages should carry a digital signature (message integrity).
NTLMSSP_NEGOTIATE_SEAL						= 0x00000020  # Messages should be encrypted (message confidentiality). 
NTLMSSP_NEGOTIATE_LM_KEY					= 0x00000080  # LM Session Key should be used for signing and sealing messages 
NTLMSSP_NEGOTIATE_NTLM						= 0x00000200  # Indicates that NTLM authentication is being used. 
NTLMSSP_ANONYMOUS							= 0x00000800  # Indicates that an anonymous context has been established. 
NTLMSSP_NEGOTIATE_OEM_DOMAIN_SUPPLIED		= 0x00001000  # Set if domain is provided in negotiate message. 
NTLMSSP_NEGOTIATE_OEM_WORKSTATION_SUPPLIED	= 0x00002000  # Set if workstation is provided in negotiate message.
NTLMSSP_NEGOTIATE_ALWAYS_SIGN				= 0x00008000  # Requests the presence of a signature block on all messages. 
NTLMSSP_TARGET_TYPE_DOMAIN					= 0x00010000  # If set, NTLM.Challenge.TargetName must be a domain name.   
NTLMSSP_TARGET_TYPE_SERVER					= 0x00020000  # If set, NTLM.Challenge.TargetName must be a server name.
NTLMSSP_NEGOTIATE_EXTENDED_SESSIONSECURITY	= 0x00080000  # Requests user of NTLMv2
NTLMSSP_NEGOTIATE_IDENTIFY					= 0x00100000  # 
NTLMSSP_REQUEST_NON_NT_SESSION_KEY			= 0x00400000  # 
NTLMSSP_NEGOTIATE_TARGET_INFO				= 0x00800000  # Requests NTLM.Challenge.TargetInfo fields be populated
NTLMSSP_NEGOTIATE_VERSION					= 0x02000000  # Requests protocol version number be provided in all messages.
NTLMSSP_NEGOTIATE_128						= 0x20000000  # 128-bit session key
NTLMSSP_NEGOTIATE_KEY_EXCH					= 0x40000000  # Request explicit key exchange. Improves security. 
NTLMSSP_NEGOTIATE_56						= 0x80000000  # 56-bit session key



class NTLM(object):
	FLAGS = NTLMSSP_NEGOTIATE_UNICODE | NTLMSSP_REQUEST_TARGET | NTLMSSP_NEGOTIATE_NTLM | NTLMSSP_NEGOTIATE_128 | NTLMSSP_NEGOTIATE_ALWAYS_SIGN | NTLMSSP_NEGOTIATE_VERSION | NTLMSSP_NEGOTIATE_EXTENDED_SESSIONSECURITY
	class Negotiate(object):
		def __init__(self):
			"""
			An NTLM.negotiate_message is sent to specify (negotiate) the clients supported NTLM options.
			"""
			self.flags = NTLM.FLAGS 
			self.message_type = 0x01
			self.domain_name_length = 0x00
			self.domain_name_maxlength = 0x00
			self.domain_name_offset = 0x28
			self.workstation_name_length = 0x00
			self.workstation_name_maxlength = 0x00
			self.workstation_name_offset = 0x28
			self.version = VERSION
			
		def encode(self):
			return pack('<8sIIHHIHHIQ', NTLMSSP_SIGNATURE, self.message_type, self.flags, self.domain_name_length, self.domain_name_maxlength, self.domain_name_offset, self.workstation_name_length, self.workstation_name_maxlength, self.workstation_name_offset, self.version)


	class Challenge(object):
		def __init__(self, token):
			"""
			An NTLM.challenge_message is received in response to the NTLM.negotiate_message
			and challenges the client to prove its identity. 
			
			token                 The NTLM token in the packet sent by the server 
			"""
			self.token = token
			self.signature, self.message_type, self.target_name_len, self.target_name_maxlength, self.target_name_offset, self.flags, self.server_challenge, self.reserved, self.target_info_length, self.target_info_maxlength, self.target_info_offset = unpack('<8sIHHII8sQHHI', self.token[0:48])
			if (self.target_name_len > 0):
				self.target_info = self.token[self.target_info_offset : self.target_info_offset + self.target_info_length]
			else:
				self.target_info = ""
		
		def encode(self):
			# Obviously this doesn't actually encode anything... This is the token received from the server
			# so it doesn't have to, it just makes the interface consistent.
			return self.token


	class Authenticate(object):
		MESSAGE_SIZE = 0x58
		def __init__(self, server_challenge, target_info, nc_message, user, password, user_domain="", server_machine_name=""):
			"""
			An NTLM.authenticate_message is sent in response to the NTLM.challenge_message to prove the clients identity.
			
			server_challenge      Challenge from the server (NTLM.challenge_message.server_challenge)
			target_info           Target info (NTLM.challenge_message.target_info)
			nc_message            Concatenation of NTLM.negotiate_message and NTLM.challenge_message
			user                  User name to authenticate, should be empty if password protected sharing is off
			password              Password to authenticate, should be empty if password protected sharing is off
			user_domain*          User's domain name (optional, unless required by server)
			server_machine_name   Server's workstation name (optional, unless required by server)
			"""
			self.flags = NTLM.FLAGS 
			self.user = user
			self.password = password
			self.user_domain = user_domain
			self.server_machine_name = server_machine_name
			self.challenge_response_lm = '\x00'*24

			if ((self.user == '') and (self.password == '')):
				# Empty user and password, for systems that have turned off password protected sharing
				self.challenge_response_nt = ''
				self.session_base_key = 0
				self.user_domain = ''
				self.workstation = ''
				self.encrypted_random_session_key = ''
				self.mic = '\x00'*16
			else:
				# Authenticate user and password
				if not(self.flags & NTLMSSP_NEGOTIATE_128):
					raise NTLM.UnsupportedFeature("Client requires 128-bit encryption but the server does not support it.")
				
				self.av_pair = AvPair(target_info)
				
				if (self.av_pair.get("MsvAvTimestamp")):
					should_supply_mic = True
					msv_av_flags = self.av_pair.get("MsvAvFlags")
					if (msv_av_flags): self.av_pair.set("MsvAvFlags", msv_av_flags | 0x02)
					else: self.av_pair.set("MsvAvFlags", 0x02)
				else:
					should_supply_mic = False

				# The following two pairs are optional and are not supported at the moment. Empty values are supplied.  
				self.av_pair.set("MsvAvChannelBindings", "\x00"*16)
				self.av_pair.set("MsvAvTargetName", "")

				# Calculate the challenge response and session key				
				self.challenge_response_nt, self.challenge_response_lm, self.session_base_key = self._get_v2_response(self.user, self.user_domain, self.password, server_challenge, target_info)
			
				# If NTLM v2 is used, the key exchange key MUST be the 128-bit session base key.
				# If supporting v1 in the future refer to [MS-NTLM] section 3.4.5.1 KXKEY, page 65 for calculations. 
				self.key_exchange_key = self.session_base_key
			
				# We are not setting the flag NTLMSSP_NEGOTIATE_KEY_EXCHANGE and thus the following keys are defined as such
				# If supporting in the future refer to [MS-NTLM] section 3.1.5.1.2 Client Receives a CHALLENGE_MESSAGE from the Server, page 46.
				self.exported_session_key = self.key_exchange_key
				self.encrypted_random_session_key = "" 
				
				#m = md5(self.exported_session_key + "session key to client-to-server signing key magic constant")
				#self.client_signing_key = m.digest()
				
				# If NTLM v2 authentication is used and the CHALLENGE_MESSAGE contains a TargetInfo field, the client
				# SHOULD NOT send the LmChallengeResponse and SHOULD set the LmChallengeResponseLen and
				# LmChallengeResponseMaxLen fields in the AUTHENTICATE_MESSAGE to zero.
				if (target_info):
					# TODO: Windows 7 sends the LM reponse as Z(24), but [MS-NLMP] says it should be empty. 
					# Windows 7 accepts both, but should test with different version servers for compatibility. 
					#self.challenge_response_lm = "\x00"*24
					self.challenge_response_lm = ""
	
				self.mic = "\x00"*16	
				if (should_supply_mic):
					self.mic = hmac.new(self.exported_session_key, nc_message + self.encode()).digest()
	
		def _add_field(self, data):
			""" Adds a field to the authenticate message with the length and offset of the data """ 
			self.message += pack('<HHI', len(data), len(data), NTLM.Authenticate.MESSAGE_SIZE + len(self.payload))
			self.payload += data

		def _get_v2_response(self, user, user_domain, password, server_challenge, target_info):
			""" Returns the LM and NTLM version 2 response messages """
			# Refer to [MS-NTLM] section 3.3.2 NTLM v2 Authentication, page 58.
			# Calculate response key
			m = md4.MD4()
			m.update(password.encode('utf-16-le'))
			password_hash = m.digest()
			response_key = hmac.new(password_hash, (user.upper() + user_domain).encode('utf-16-le')).digest()
			
			# Look for MsvAvTimestamp in target_info. Use it if it exists, otherwise use local system time.
			timestamp = self.av_pair.get("MsvAvTimestamp")
			if not(timestamp): 
				timestamp = pack('>Q', int((time.time() + 11644473600) * 10000000))
				self.av_pair.set("MsvAvTimestamp", timestamp)

			client_challenge = pack('<Q', random.randint(0, pow(2, 64) - 1))
			
			# Calculate responses and session base key
			temp = "\x01"*2 + "\x00"*6 + timestamp + client_challenge + "\x00"*4 + self.av_pair.encode()
			proof_str_nt = hmac.new(response_key, server_challenge + temp).digest()
			challenge_response_nt = proof_str_nt + temp
			challenge_response_lm = hmac.new(response_key, server_challenge + client_challenge).digest() + client_challenge
			session_base_key = hmac.new(response_key, proof_str_nt).digest()
			
			return challenge_response_nt, challenge_response_lm, session_base_key

		def _pack_16(self, endianess, in_16):
			""" Pack a 16-byte integer into a string """
			lo_8 = (in_16 & pow(2,64) - 1)
			hi_8 = ((in_16 & (pow(2,64) - 1) << 64) >> 64)
			if (endianess == "<"):
				return "%s%s" % (pack('<Q', lo_8), pack('<Q', hi_8))
			else:
				return "%s%s" % (pack('>Q', hi_8), pack('>Q', lo_8))
		
		def encode(self):
			""" Get an authenticate packet """
			self.message = 'NTLMSSP\x00\x03\x00\x00\x00' # NTLM signature and authenticate message type
			self.version = VERSION
			self.payload = ""
			
			self._add_field(self.challenge_response_lm)
			self._add_field(self.challenge_response_nt)
			self._add_field(self.user_domain.encode("utf-16-le"))
			self._add_field(self.user.encode("utf-16-le"))
			self._add_field(self.server_machine_name.encode("utf-16-le"))
			self._add_field(self.encrypted_random_session_key)

			self.message += pack('<IQ16s', self.flags, self.version, self.mic)
			
			return self.message + self.payload 

	class UnsupportedFeature(Exception): pass


class AvPair(object):
	def __init__(self, target_info=None):
		self.av_id = ["MsvAvEOL", "MsvAvNbComputerName", "MsvAvNbDomainName", "MsvAvDnsComputerName", "MsvAvDnsDomainName", "MsvAvDnsTreeName", "MsvAvFlags", "MsvAvTimestamp", "MsAvRestrictions", "MsvAvTargetName", "MsvAvChannelBindings"]
		if (target_info): self.parse(target_info)
		else: self._new()
		
	def __repr__(self):
		return self.av_pair.__repr__()
	
	def _new(self):
		self.av_pair = {}
	
	def parse(self, target_info):
		self._new()
		offset = 0
		while (len(target_info) - offset >= 4):
			id, length = unpack('<HH', target_info[0 + offset : 4 + offset])
			if (id == 0): break
			else: self.av_pair[self.av_id[id]] = target_info[4 + offset : 4 + length + offset]
			offset += (4 + length)
		
	def get(self, id):
		try:
			if (id.__class__.__name__ == "str"): return self.av_pair[id]
			else: return self.av_pair[self.av_id[id]]
		except:
			return None

	def set(self, id, value):
		if (id.__class__.__name__ == "str"): self.av_pair[id] = value
		else: self.av_pair[self.av_id[id]] = value

	def encode(self):
		target_info = ""
		for id in xrange(1, len(self.av_id)):
			key = self.av_id[id]
			if (self.av_pair.has_key(key)):
				value = self.av_pair[key]
				if (key == "MsvAvFlags"): value = pack("<I", value)
				target_info += pack('<HH', self.av_id.index(key), len(value)) + value
		target_info += "\x00"*4
		
		return target_info
	

class GSS(object):
	STATUS = {0:(0,"ACCEPT_COMPLETED"), 1:(1,"ACCEPT_INCOMPLETE"), 2:(2,"REJECT"), 3:(3,"REQUEST_MIC")}

	@classmethod
	def ber_write(cls, dlength):
		if (dlength <= 0): raise Exception('dlength must be greater than 0')
		if (dlength > 127):
			dlength += int(math.pow(256, int(math.floor(math.log(dlength, 256))+1))*(128+int(math.floor(math.log(dlength, 256))+1)))
		size = int(math.floor(math.log(dlength, 256))+1)
		data = ''
		for i in range(size,0,-1):
			data += pack('B',dlength/math.pow(256,i-1))
			dlength %= math.pow(256,i-1)
		return data

	@classmethod
	def ber_read(cls, data):
		if (len(data) == 0): raise Exception('Data length too short')
		first = unpack('<B',data[0])[0]
		num = 0
		if (first&128): # bit 8 on first octet indicates multiple octets required to store the integer
			length = first^128 # bits 7-0 indicate the number of octets that follow
			if (len(data) < length): raise Exception('Data length too short')
			for i in range(length):
				num *= 256
				num += unpack('<B', data[i+1])[0]
			return num, length+1
		return first, 1
	
	@staticmethod
	def _get_field(data):
		id1 = data[0]
		_, bytes = GSS.ber_read(data[1:])
		data = data[bytes+1:]
		id2 = data[0]
		length, bytes = GSS.ber_read(data[1:])
		id = (ord(id1)*256+ord(id2))
		value = data[bytes+1:bytes+1+length]
		remaining = data[bytes+1+length:]
		return id, value, remaining
	
	@staticmethod
	def neg_token_init(ntlm_message):
		mech_token = '\x04' + GSS.ber_write(len(ntlm_message)) + ntlm_message
		mech_token = '\xa2' + GSS.ber_write(len(mech_token)) + mech_token

		mech_type_list = '\x06' + GSS.ber_write(len(NTLMSSP_OID)) + NTLMSSP_OID
		mech_type_list = '\x30' + GSS.ber_write(len(mech_type_list)) + mech_type_list
		mech_type_list = '\xa0' + GSS.ber_write(len(mech_type_list)) + mech_type_list

		neg_token = mech_type_list + mech_token
		neg_token = '\x30' + GSS.ber_write(len(neg_token)) + neg_token
		neg_token = '\xa0' + GSS.ber_write(len(neg_token)) + neg_token
		neg_token = '\x06' + GSS.ber_write(len(SPNEGO_OID)) + SPNEGO_OID + neg_token
		neg_token = '\x60' + GSS.ber_write(len(neg_token)) + neg_token

		return neg_token
	
	@staticmethod
	def unwrap_neg_token_resp(gss_token):	
		ntlm_token = None
		id, neg_token_resp, _ = GSS._get_field(gss_token)
		while (len(neg_token_resp) > 0):
			id, value, neg_token_resp = GSS._get_field(neg_token_resp)
			if (id == 0xa00a): neg_state = GSS.STATUS[ord(value)]
			if (id == 0xa106): pass # TODO: supportedMech, will assume NTLMSSP is supported for now but fix this in the future
			if (id == 0xa204): ntlm_token = value
			if (id == 0xa304): pass # mechListMIC, not used for now
		
		return neg_state, ntlm_token 

	@staticmethod
	def neg_token_resp(ntlm_message):
		response_token = '\x04' + GSS.ber_write(len(ntlm_message)) + ntlm_message
		response_token = '\xa2' + GSS.ber_write(len(response_token)) + response_token
		
		neg_state = '\xa0\x03\x0a\x01\x01'
		neg_token =  neg_state + response_token
		neg_token = '\x30' + GSS.ber_write(len(neg_token)) + neg_token
		neg_token = '\xa1' + GSS.ber_write(len(neg_token)) + neg_token
		
		return neg_token


